"""
Registry 系统设计的核心原则
1. 类型保真原则 (Type Fidelity):注册系统不应该改变或丢失组件的原始类型信息。
2. 里氏替换原则 (LSP): ex.虽然 BaseHookFactory 是 HookFactoryInterface 的子类，但 Registry 应该能够返回具体类型，而不仅仅是接口类型。
3. 双重类型契约 (Dual Type Contract)
    Registry 应该同时保证：
        接口契约：所有注册的类都实现了必需的接口
        具体类型保留：能够返回原始的具体类型
4. 透明性原则 (Transparency): Registry 应该像一个透明的容器，不改变存储对象的性质。
5. 灵活性原则 (Flexibility): 提供多种访问方式，满足不同的使用场景。

# 这在 Python 中是不可能的
# Python 类型系统的根本限制是：无法建立字符串到类型的静态映射。
registry.get("resnet") -> type[ResNet]  # 自动推断
registry.get("vgg") -> type[VGG]       # 自动推断
"""

from __future__ import annotations

import importlib
import inspect
import logging
from collections.abc import Callable, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import (
    Any,
    ClassVar,
    Generic,
    Protocol,
    TypeVar,
    overload,
    runtime_checkable,
)

from continuallearning.interfaces import ComponentInterface

logger = logging.getLogger(__name__)

# Type variables
# I: Interface constraint - ensures all components implement required interface
I = TypeVar("I", bound=ComponentInterface)
# T: Concrete type - preserves the actual registered class type
T = TypeVar("T")
# C: Class decorator type - preserves type through decoration
C = TypeVar("C")


@dataclass(frozen=True)
class RegistryInfo:
    """Immutable metadata about a registered component."""

    name: str
    registries: frozenset[str] = field(default_factory=frozenset)

    def with_registry(self, registry_name: str) -> RegistryInfo:
        """Create a new RegistryInfo with an additional registry."""
        return RegistryInfo(
            name=self.name, registries=self.registries | {registry_name}
        )


@runtime_checkable
class RegistryAware(Protocol):
    """Protocol for components that track their registry metadata."""

    __registry_info__: ClassVar[RegistryInfo]


class RegistrationError(Exception):
    """Raised when component registration fails."""


class ComponentNotFoundError(KeyError):
    """Raised when a requested component is not found in the registry."""

    def __init__(self, name: str, registry_name: str, available: list[str]):
        self.name = name
        self.registry_name = registry_name
        self.available = available
        super().__init__(
            f"{name} is not registered in {registry_name} registry. "
            f"Available components: {available}"
        )


class Registry(Generic[I]):
    """
    Type-preserving registry for managing components implementing a specific interface.

    The Registry provides a central location to register components while preserving
    their concrete types. It ensures type safety at both the interface level and
    the concrete implementation level.

    Key Features:
    - Dual type contract: Interface compliance + concrete type preservation
    - Type-aware retrieval methods for maximum flexibility
    - Case-insensitive name lookups with preserved original casing
    - Automatic component discovery from packages
    - Thread-safe operations with context managers
    - Full type transparency - registry doesn't alter component types

    Example:
        >>> # Define interface and implementations
        >>> class ModelInterface(ComponentInterface):
        ...     def forward(self, x): ...
        >>>
        >>> class ResNet(ModelInterface):
        ...     def forward(self, x): return x
        ...     def resnet_specific_method(self): return "ResNet"
        >>>
        >>> # Create registry and register component
        >>> registry = Registry("Model", ModelInterface)
        >>> @registry.register("resnet")
        ... class ResNet(ModelInterface): ...
        >>>
        >>> # Type-preserving retrieval
        >>> ResNetClass = registry.get_typed("resnet", ResNet)
        >>> model = ResNetClass()  # Full ResNet type available
        >>> model.resnet_specific_method()  # Works!
    """

    def __init__(self, name: str, interface_type: type[I]) -> None:
        """
        Initialize a new registry.

        Args:
            name: Human-readable name for this registry
            interface_type: Interface type that components must implement

        Raises:
            TypeError: If interface_type is not a valid type
        """
        if not isinstance(interface_type, type):
            raise TypeError(
                f"interface_type must be a valid type, got {type(interface_type).__name__}"
            )

        self.name = name
        self._interface_type = interface_type
        # Store components as Any to preserve concrete types
        self._components: dict[str, type[Any]] = {}
        self._original_names: dict[str, str] = {}
        self._locked = False

    def __len__(self) -> int:
        """Return the number of registered components."""
        return len(self._components)

    def __contains__(self, key: str) -> bool:
        """Check if a component name is registered (case-insensitive)."""
        return key.lower() in self._components

    def __repr__(self) -> str:
        """Return a string representation of the registry."""
        names = sorted(self._original_names.values())
        return f"<{self.name} Registry: {len(self)} components {names}>"

    @contextmanager
    def locked(self):
        """Context manager to temporarily lock the registry from modifications."""
        old_locked = self._locked
        self._locked = True
        try:
            yield
        finally:
            self._locked = old_locked

    def _validate_component(self, cls: type) -> None:
        """
        Validate that a class implements the required interface.

        Args:
            cls: Class to validate

        Raises:
            RegistrationError: If validation fails
        """
        if not inspect.isclass(cls):
            raise RegistrationError(f"Expected a class, got {type(cls).__name__}")

        if not issubclass(cls, self._interface_type):
            raise RegistrationError(
                f"Class {cls.__name__} must implement {self._interface_type.__name__}"
            )

    def _update_registry_info(self, cls: type, name: str) -> None:
        """Update or create registry metadata on the class."""
        if hasattr(cls, "__registry_info__"):
            # Update existing metadata
            old_info = getattr(cls, "__registry_info__")
            new_info = old_info.with_registry(self.name)
        else:
            # Create new metadata
            new_info = RegistryInfo(name=name, registries=frozenset([self.name]))

        # Set the attribute on the class
        setattr(cls, "__registry_info__", new_info)

    # Overloaded signatures for different usage patterns
    @overload
    def register(self, name_or_cls: str) -> Callable[[type[C]], type[C]]:
        """Register with explicit name: @registry.register("name")"""
        ...

    @overload
    def register(self, name_or_cls: type[C]) -> type[C]:
        """Register without name: @registry.register"""
        ...

    def register(
        self, name_or_cls: str | type[C] | None = None
    ) -> Callable[[type[C]], type[C]] | type[C]:
        """
        Register a component with the registry while preserving its type.

        This decorator ensures that the decorated class maintains its original
        type information for both static type checkers and runtime usage.

        Can be used as a decorator with or without arguments:

            @registry.register
            class MyComponent(Interface): ...

            @registry.register("custom_name")
            class MyComponent(Interface): ...

        Args:
            name_or_cls: Component name (str) or class to register

        Returns:
            Decorator function that preserves class type, or the registered class

        Raises:
            RegistrationError: If registration fails or registry is locked
        """
        if self._locked:
            raise RegistrationError("Registry is locked and cannot be modified")

        # Handle direct class decoration: @register
        if inspect.isclass(name_or_cls):
            # Direct decoration without parentheses
            cls = name_or_cls
            self._register_class(cls, cls.__name__)
            return cls

        # Handle named decoration: @register("name") or @register()
        def decorator(cls: type[C]) -> type[C]:
            # Determine the registration name
            if isinstance(name_or_cls, str):
                name = name_or_cls
            else:
                name = cls.__name__

            self._register_class(cls, name)
            return cls  # Return the original class unchanged

        return decorator

    def _register_class(self, cls: type, name: str) -> None:
        """
        Internal method to register a class.

        Args:
            cls: Class to register
            name: Registration name

        Raises:
            RegistrationError: If validation fails
        """
        self._validate_component(cls)

        key = name.lower()

        if key in self._components:
            logger.warning(
                f"Overwriting '{name}' in {self.name} registry "
                f"(was {self._components[key].__name__}, now {cls.__name__})"
            )

        # Store as Any to preserve concrete type
        self._components[key] = cls
        self._original_names[key] = name
        self._update_registry_info(cls, name)

        logger.debug(f"Registered {cls.__name__} as '{name}' in {self.name} registry")

    def unregister(self, name: str) -> type[Any] | None:
        """
        Remove a component from the registry.

        Args:
            name: Name of component to remove

        Returns:
            The removed component class with preserved type, or None if not found

        Raises:
            RegistrationError: If registry is locked
        """
        if self._locked:
            raise RegistrationError("Registry is locked and cannot be modified")

        key = name.lower()
        cls = self._components.pop(key, None)
        self._original_names.pop(key, None)

        if cls:
            logger.debug(f"Unregistered '{name}' from {self.name} registry")

        return cls

    def get(self, name: str) -> type[Any]:
        """
        Get a registered component class.

        Returns the actual registered type. The return type is `type[Any]`
        because Python's type system cannot statically determine the concrete
        type from a runtime string parameter.

        For better type safety, consider:
        1. Using get_typed() with expected type
        2. Manually casting the result if you know the type
        """
        key = name.lower()

        if key not in self._components:
            available = sorted(self._original_names.values())
            raise ComponentNotFoundError(name, self.name, available)

        return self._components[key]

    @overload
    def get_typed(self, name: str, expected_type: type[T]) -> type[T]:
        """Get with explicit type hint."""
        ...

    @overload
    def get_typed(self, name: str, expected_type: None = None) -> type[Any]:
        """Get with preserved concrete type."""
        ...

    def get_typed(
        self, name: str, expected_type: type[T] | None = None
    ) -> type[T] | type[Any]:
        """
        Get a component class with preserved concrete type.

        This method preserves the original type information, allowing
        access to implementation-specific methods and attributes.

        Args:
            name: Component name (case-insensitive)
            expected_type: Optional type hint for better IDE support

        Returns:
            The registered component class with original type preserved

        Raises:
            ComponentNotFoundError: If component not found
            TypeError: If expected_type provided and doesn't match

        Example:
            >>> ResNetClass = registry.get_typed("resnet", ResNet)
            >>> # or without type hint (still preserves type at runtime)
            >>> ResNetClass = registry.get_typed("resnet")
        """
        key = name.lower()

        if key not in self._components:
            available = sorted(self._original_names.values())
            raise ComponentNotFoundError(name, self.name, available)

        cls = self._components[key]

        if expected_type is not None and not issubclass(cls, expected_type):
            raise TypeError(
                f"Component '{name}' is {cls.__name__}, "
                f"not compatible with expected type {expected_type.__name__}"
            )

        return cls

    def create(self, name: str, **kwargs: Any) -> I:
        """
        Create a component instance by name (interface typed).

        Args:
            name: Component name
            **kwargs: Constructor arguments

        Returns:
            Component instance typed as interface

        Raises:
            ComponentNotFoundError: If component not found
            TypeError: If instantiation fails
        """
        component_cls = self.get(name)
        return self._instantiate(component_cls, name, kwargs)

    @overload
    def create_typed(self, name: str, expected_type: type[T], **kwargs: Any) -> T:
        """Create with explicit type."""
        ...

    @overload
    def create_typed(self, name: str, expected_type: None = None, **kwargs: Any) -> Any:
        """Create with preserved type."""
        ...

    def create_typed(
        self, name: str, expected_type: type[T] | None = None, **kwargs: Any
    ) -> T | Any:
        """
        Create a component instance with preserved concrete type.

        This allows accessing implementation-specific methods on the created instance.

        Args:
            name: Component name
            expected_type: Optional type hint for better IDE support
            **kwargs: Constructor arguments

        Returns:
            Component instance with concrete type preserved

        Raises:
            ComponentNotFoundError: If component not found
            TypeError: If instantiation fails or type doesn't match

        Example:
            >>> model = registry.create_typed("resnet", ResNet, layers=50)
            >>> model.resnet_specific_method()  # Works!
        """
        component_cls = self.get_typed(name, expected_type)
        return self._instantiate(component_cls, name, kwargs)

    def _instantiate(
        self, component_cls: type[T], name: str, kwargs: dict[str, Any]
    ) -> T:
        """
        Internal method to instantiate a component.

        Args:
            component_cls: Component class to instantiate
            name: Component name (for error messages)
            kwargs: Constructor arguments

        Returns:
            Component instance

        Raises:
            TypeError: If instantiation fails
        """
        try:
            # Filter kwargs based on constructor signature
            if hasattr(component_cls, "__init__"):
                sig = inspect.signature(component_cls.__init__)

                # Check if constructor accepts **kwargs
                accepts_kwargs = any(
                    p.kind == p.VAR_KEYWORD for p in sig.parameters.values()
                )

                if not accepts_kwargs:
                    # Filter to valid parameter names only
                    valid_params = {
                        p.name
                        for p in sig.parameters.values()
                        if p.name != "self" and p.kind != p.VAR_KEYWORD
                    }
                    kwargs = {k: v for k, v in kwargs.items() if k in valid_params}

            return component_cls(**kwargs)

        except Exception as e:
            logger.error(f"Failed to create {self.name} '{name}': {e}", exc_info=True)
            raise TypeError(f"Failed to instantiate {self.name} '{name}': {e}") from e

    def build(self, config: Mapping[str, Any]) -> I:
        """
        Build a component from configuration (interface typed).

        Args:
            config: Configuration dict with 'type' key

        Returns:
            Component instance typed as interface

        Raises:
            TypeError: If config invalid
            KeyError: If 'type' missing
        """
        if not isinstance(config, Mapping):
            raise TypeError(f"Config must be a mapping, got {type(config).__name__}")

        if "type" not in config:
            raise KeyError("Configuration must contain 'type' key")

        # Create mutable copy
        config_dict = dict(config)
        component_type = config_dict.pop("type")

        return self.create(component_type, **config_dict)

    def build_typed(
        self, config: Mapping[str, Any], expected_type: type[T] | None = None
    ) -> T | Any:
        """
        Build a component from configuration with preserved type.

        Args:
            config: Configuration dict with 'type' key
            expected_type: Optional type hint

        Returns:
            Component instance with preserved type

        Raises:
            TypeError: If config invalid or type doesn't match
            KeyError: If 'type' missing
        """
        if not isinstance(config, Mapping):
            raise TypeError(f"Config must be a mapping, got {type(config).__name__}")

        if "type" not in config:
            raise KeyError("Configuration must contain 'type' key")

        # Create mutable copy
        config_dict = dict(config)
        component_type = config_dict.pop("type")

        return self.create_typed(component_type, expected_type, **config_dict)

    def list_registered(self) -> list[str]:
        """
        List all registered component names.

        Returns:
            Sorted list of original component names
        """
        return sorted(self._original_names.values())

    def list_components(self) -> dict[str, type[Any]]:
        """
        Get all registered components with their concrete types.

        Returns:
            Dictionary mapping names to component classes with preserved types
        """
        return {self._original_names[key]: cls for key, cls in self._components.items()}

    def discover_components(
        self, package_paths: list[str], auto_register_attr: str = "__auto_register__"
    ) -> int:
        """
        Auto-discover and register components from packages.

        Args:
            package_paths: Package paths to scan
            auto_register_attr: Attribute name to check for auto-registration

        Returns:
            Number of components discovered

        Example:
            >>> count = registry.discover_components(['mypackage.models'])
            >>> print(f"Discovered {count} components")
        """
        discovered = 0

        for package_path in package_paths:
            try:
                discovered += self._discover_in_package(
                    package_path, auto_register_attr
                )
            except Exception as e:
                logger.warning(f"Failed to discover components in {package_path}: {e}")

        return discovered

    def _discover_in_package(self, package_path: str, auto_register_attr: str) -> int:
        """Discover components in a single package."""
        package = importlib.import_module(package_path)

        if not hasattr(package, "__path__"):
            # Single module, not a package
            return self._discover_in_module(package, auto_register_attr)

        package_file = package.__file__
        if package_file is None:
            logger.warning(f"Cannot determine file path for package {package_path}")
            return 0

        package_dir = Path(package_file).parent
        count = 0

        # Scan all Python files
        for py_file in package_dir.glob("**/*.py"):
            if py_file.name.startswith("_"):
                continue

            # Convert file path to module path
            relative_path = py_file.relative_to(package_dir)
            module_parts = list(relative_path.parts[:-1]) + [relative_path.stem]
            module_path = f"{package_path}.{'.'.join(module_parts)}"

            try:
                module = importlib.import_module(module_path)
                count += self._discover_in_module(module, auto_register_attr)
            except Exception as e:
                logger.debug(f"Skipping {module_path}: {e}")

        return count

    def _discover_in_module(self, module: Any, auto_register_attr: str) -> int:
        """Discover components in a single module."""
        count = 0

        for attr_name in dir(module):
            if attr_name.startswith("_"):
                continue

            try:
                attr = getattr(module, attr_name)

                if (
                    inspect.isclass(attr)
                    and issubclass(attr, self._interface_type)
                    and attr is not self._interface_type
                    and getattr(attr, auto_register_attr, False)
                ):
                    self.register(attr)
                    count += 1
                    logger.info(
                        f"Auto-registered {attr.__name__} from {module.__name__}"
                    )
            except Exception as e:
                logger.debug(f"Error checking {attr_name}: {e}")

        return count


# Helper functions for creating typed registries
def create_registry(name: str, interface_type: type[I]) -> Registry[I]:
    """
    Create a new registry with proper typing.

    Args:
        name: Registry name
        interface_type: Interface type

    Returns:
        Typed registry instance
    """
    return Registry(name, interface_type)


# Type alias for better readability
TypePreservingRegistry = Registry
